package org.finesys.chat.core.base.embedding.store.milvus;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingMatch;
import org.finesys.chat.core.base.embedding.EmbeddingSearchRequest;
import org.finesys.chat.core.base.embedding.EmbeddingSearchResult;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.milvus.filter.MilvusMetadataFilterMapper;
import org.finesys.chat.core.base.embedding.store.util.Generator;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;

import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;

public class MilvusEmbeddingStore implements EmbeddingStore<TextSegment>, AutoCloseable {
    private final MilvusServiceClient milvusClient;
    private final String collectionName;
    private final MetricType metricType;
    private final ConsistencyLevelEnum consistencyLevel;
    private final boolean retrieveEmbeddingsOnSearch;
    private final boolean autoFlushOnInsert;
    private final FieldDefinition fieldDefinition;


    public MilvusEmbeddingStore(String host, Integer port, String collectionName, Integer dimension, IndexType indexType, MetricType metricType, String uri, String token, String username, String password, ConsistencyLevelEnum consistencyLevel, Boolean retrieveEmbeddingsOnSearch, Boolean autoFlushOnInsert, String databaseName, String idFieldName, String textFieldName, String metadataFiledName, String vectorFiledName) {
        this(createMilvusClient(host, port, uri, token, username, password, databaseName), collectionName, dimension, indexType, metricType, consistencyLevel, retrieveEmbeddingsOnSearch, autoFlushOnInsert, idFieldName, textFieldName, metadataFiledName, vectorFiledName);
    }


    public MilvusEmbeddingStore(MilvusServiceClient milvusClient, String collectionName, Integer dimension, IndexType indexType, MetricType metricType, ConsistencyLevelEnum consistencyLevel, Boolean retrieveEmbeddingsOnSearch, Boolean autoFlushOnInsert, String idFieldName, String textFieldName, String metadataFiledName, String vectorFiledName) {
        this.milvusClient = (MilvusServiceClient) ValidationUtil.ensureNotNull(milvusClient, "milvusClient");
        this.collectionName = (String) Utils.getOrDefault(collectionName, "default");
        this.metricType = (MetricType) Utils.getOrDefault(metricType, MetricType.COSINE);
        this.consistencyLevel = (ConsistencyLevelEnum) Utils.getOrDefault(consistencyLevel, ConsistencyLevelEnum.EVENTUALLY);
        this.retrieveEmbeddingsOnSearch = (Boolean) Utils.getOrDefault(retrieveEmbeddingsOnSearch, false);
        this.autoFlushOnInsert = (Boolean) Utils.getOrDefault(autoFlushOnInsert, false);
        this.fieldDefinition = new FieldDefinition((String) Utils.getOrDefault(idFieldName, "id"), (String) Utils.getOrDefault(textFieldName, "text"), (String) Utils.getOrDefault(metadataFiledName, "metadata"), (String) Utils.getOrDefault(vectorFiledName, "vector"));
        if (!CollectionOperationsExecutor.hasCollection(this.milvusClient, this.collectionName)) {
            CollectionOperationsExecutor.createCollection(this.milvusClient, this.collectionName, this.fieldDefinition, (Integer) ValidationUtil.ensureNotNull(dimension, "dimension"));
            CollectionOperationsExecutor.createIndex(this.milvusClient, this.collectionName, this.fieldDefinition.getVectorFieldName(), (IndexType) Utils.getOrDefault(indexType, IndexType.FLAT), this.metricType);
        }

        CollectionOperationsExecutor.loadCollectionInMemory(this.milvusClient, collectionName);
    }

    private static MilvusServiceClient createMilvusClient(String host, Integer port, String uri, String token, String username, String password, String databaseName) {
        ConnectParam.Builder connectBuilder = ConnectParam.newBuilder()
                .withHost((String) Utils.getOrDefault(host, "localhost"))
                .withPort((Integer) Utils.getOrDefault(port, 19530))
                .withUri(uri)
                .withToken(token)
                .withAuthorization((String) Utils.getOrDefault(username, ""), (String) Utils.getOrDefault(password, ""));
        if (databaseName != null) {
            connectBuilder.withDatabaseName(databaseName);
        }
        ConnectParam connectParam = connectBuilder.build();
        return new MilvusServiceClient(connectParam);
    }

    @Override
    public String add(Embedding embedding) {
        String id = UUID.randomUUID().toString();
        this.add(embedding, id);
        return id;
    }

    @Override
    public String add(Embedding embedding, String key) {
        this.addInternal(key, embedding, null);
        return key;
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        String id = UUID.randomUUID().toString();
        this.addInternal(id, embedding, textSegment);
        return id;
    }


    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> ids = Generator.generateRandomIds(embeddings.size());
        return this.addAll(ids, embeddings, null);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        List<String> ids = Generator.generateRandomIds(embeddings.size());
        return this.addAll(ids, embeddings, textSegments);
    }

    @Override
    public List<String> addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> textSegments) {
        ValidationUtil.ensureNotEmpty(ids, "ids");
        ValidationUtil.ensureNotEmpty(embeddings, "embeddings");
        if (ids.size() != embeddings.size()) {
            throw new IllegalArgumentException("ids and embeddings size must be equal");
        }
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field(this.fieldDefinition.getIdFieldName(), ids));
        fields.add(new InsertParam.Field(this.fieldDefinition.getTextFieldName(), Mapper.toText(textSegments, ids.size())));
        fields.add(new InsertParam.Field(this.fieldDefinition.getVectorFieldName(), Mapper.toVector(embeddings)));
        fields.add(new InsertParam.Field(this.fieldDefinition.getMetadataFieldName(), Mapper.toMetadataJsons(textSegments, ids.size())));
        CollectionOperationsExecutor.insert(this.milvusClient, this.collectionName, fields);
        if (this.autoFlushOnInsert) {
            CollectionOperationsExecutor.flush(this.milvusClient, this.collectionName);
        }
        return ids;
    }

    private void addInternal(String id, Embedding embedding, TextSegment textSegment) {
        this.addAll(Collections.singletonList(id), Collections.singletonList(embedding), textSegment == null ? null : Collections.singletonList(textSegment));
    }

    @Override
    public void delete(String key) {
        deleteAll(Collections.singletonList(key));
    }

    @Override
    public void deleteAll(List<String> keys) {
        ValidationUtil.ensureNotEmpty(keys, "keys");
        String expr = String.format("%s IN [%s]", this.fieldDefinition.getIdFieldName(), keys);
        CollectionOperationsExecutor.removeForVector(milvusClient, collectionName, expr);
    }

    @Override
    public void deleteAll(Filter filter) {
        ValidationUtil.ensureNotNull(filter, "filter");
        CollectionOperationsExecutor.removeForVector(milvusClient, collectionName, MilvusMetadataFilterMapper.map(filter, this.fieldDefinition.getMetadataFieldName()));
    }

    @Override
    public void deleteAll() {
        String expr = String.format("%s !=\"\"", this.fieldDefinition.getIdFieldName());
        CollectionOperationsExecutor.removeForVector(milvusClient, collectionName, expr);
    }

    @Override
    public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest request) {
        SearchParam searchParam = CollectionRequestBuilder.buildSearchRequest(this.collectionName, this.fieldDefinition
                , request.getQueryEmbedding().vectorAsList(), request.getFilter(), request.getLimit(), this.metricType, this.consistencyLevel);
        SearchResultsWrapper searchResultsWrapper = CollectionOperationsExecutor.search(milvusClient, searchParam);
        List<EmbeddingMatch<TextSegment>> matches = Mapper.toEmbeddingMatches(this.milvusClient, searchResultsWrapper, this.collectionName, this.fieldDefinition, this.consistencyLevel, this.retrieveEmbeddingsOnSearch);
        List<EmbeddingMatch<TextSegment>> result = matches.stream().filter(match -> match.getScore() >= request.getMinScore()).collect(Collectors.toList());
        return new EmbeddingSearchResult<>(result);
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private MilvusServiceClient milvusClient;
        private String host;
        private Integer port;
        private String collectionName;
        private Integer dimension;
        private IndexType indexType;
        private MetricType metricType;
        private String uri;
        private String token;
        private String username;
        private String password;
        private ConsistencyLevelEnum consistencyLevel;
        private Boolean retrieveEmbeddingsOnSearch;
        private String databaseName;
        private Boolean autoFlushOnInsert;
        private String idFieldName;
        private String textFieldName;
        private String metadataFieldName;
        private String vectorFieldName;

        public Builder() {
        }

        public Builder milvusClient(MilvusServiceClient milvusClient) {
            this.milvusClient = milvusClient;
            return this;
        }

        public Builder host(String host) {
            this.host = host;
            return this;
        }

        public Builder port(Integer port) {
            this.port = port;
            return this;
        }

        public Builder collectionName(String collectionName) {
            this.collectionName = collectionName;
            return this;
        }

        public Builder dimension(Integer dimension) {
            this.dimension = dimension;
            return this;
        }

        public Builder indexType(IndexType indexType) {
            this.indexType = indexType;
            return this;
        }

        public Builder metricType(MetricType metricType) {
            this.metricType = metricType;
            return this;
        }

        public Builder uri(String uri) {
            this.uri = uri;
            return this;
        }

        public Builder token(String token) {
            this.token = token;
            return this;
        }

        public Builder username(String username) {
            this.username = username;
            return this;
        }

        public Builder password(String password) {
            this.password = password;
            return this;
        }

        public Builder consistencyLevel(ConsistencyLevelEnum consistencyLevel) {
            this.consistencyLevel = consistencyLevel;
            return this;
        }

        public Builder retrieveEmbeddingsOnSearch(Boolean retrieveEmbeddingsOnSearch) {
            this.retrieveEmbeddingsOnSearch = retrieveEmbeddingsOnSearch;
            return this;
        }

        public Builder autoFlushOnInsert(Boolean autoFlushOnInsert) {
            this.autoFlushOnInsert = autoFlushOnInsert;
            return this;
        }

        public Builder databaseName(String databaseName) {
            this.databaseName = databaseName;
            return this;
        }

        public Builder idFieldName(String idFieldName) {
            this.idFieldName = idFieldName;
            return this;
        }

        public Builder textFieldName(String textFieldName) {
            this.textFieldName = textFieldName;
            return this;
        }

        public Builder metadataFieldName(String metadataFieldName) {
            this.metadataFieldName = metadataFieldName;
            return this;
        }

        public Builder vectorFieldName(String vectorFieldName) {
            this.vectorFieldName = vectorFieldName;
            return this;
        }

        public MilvusEmbeddingStore build() {
            return this.milvusClient == null ? new MilvusEmbeddingStore(this.host, this.port, this.collectionName, this.dimension, this.indexType, this.metricType, this.uri, this.token, this.username, this.password, this.consistencyLevel, this.retrieveEmbeddingsOnSearch, this.autoFlushOnInsert, this.databaseName, this.idFieldName, this.textFieldName, this.metadataFieldName, this.vectorFieldName) : new MilvusEmbeddingStore(this.milvusClient, this.collectionName, this.dimension, this.indexType, this.metricType, this.consistencyLevel, this.retrieveEmbeddingsOnSearch, this.autoFlushOnInsert, this.idFieldName, this.textFieldName, this.metadataFieldName, this.vectorFieldName);
        }
    }

    @Override
    public void close() {
        if (milvusClient != null) {
            this.milvusClient.close();
        }
    }

    @Override
    public void stop() {
        this.close();
    }
}
